package com.leetcode.algs4.search;

/**
 * @author Dennis Li
 * @date 2020/6/15 19:43
 */
public class RedBlackTree<Key extends Comparable, Value> {
    private static final boolean RED = true;
    private static final boolean BLACK = false;

    private Node root;

    // BST helper node data type
    private class Node {
        private Key key;           // key
        private Value val;         // associated data
        private Node left, right;  // links to left and right subtrees
        private boolean color;     // color of parent link
        private int size;          // subtree count

        public Node(Key key, Value val, boolean color, int size) {
            this.key = key;
            this.val = val;
            this.color = color;
            this.size = size;
        }
    }

    private void put(Key key, Value val) {
        root = put(root, key, val);
        root.color = BLACK;
    }

    private Node put(Node h, Key key, Value val) {
        // 如果没有则插入一个红色结点，这样方便以后调整
        if (h == null)
            return new Node(key, val, RED, 1);

        int cmp = key.compareTo(h.key);

        // 满足BST的性质，如果小于，则从左边插入
        if (cmp < 0) h.left = put(h.left, key, val);
        else if (cmp > 0) h.right = put(h.right, key, val);
        // 已经在树中，更新值
        else h.val = val;

        /**
         * 红黑树调整的目的是为了在插入结点的同时动态调整树的高度
         * 避免树的高度过高，提升索引的效率
         */

        // 重点！！ 红黑树的调整过程
        // 左黑右红
        if (!isRed(h.left) && isRed(h.right))
            h = rotateLeft(h); // 左旋
        // 左红右黑
        if (isRed(h.left) && !isRed(h.right))
            h = rotateRight(h); // 右旋
        // 全红
        if (isRed(h.left) && isRed(h.right))
            flipColors(h); // 颜色反转 -- 结点成为2-结点，父节点成为3-结点

        // 调整树的规模
        h.size = size(h.left) + size(h.right) + 1;
        return h;
    }

    private boolean isRed(Node x) {
        if (x == null) return false;
        return x.color == RED;
    }

    private int size(Node x) {
        if (x == null) return 0;
        return x.size;
    }

    private Node rotateLeft(Node h) {
        Node x = h.right;
        x.left = h;
        x.color = h.color;
        // 与父节点链接变为红色
        h.color = RED;
        x.size = size(h.left) + size(h.right) + 1;
        return x;
    }

    private Node rotateRight(Node h) {
        Node x = h.left;
        h.left = x.right;
        x.right = h;
        x.color = h.color;
        h.color = RED;
        x.size = size(h.left) + size(h.right) + 1;
        return x;
    }

    private void flipColors(Node h) {
        if (h == null) return;
        h.color = RED;
        h.left.color = BLACK;
        h.right.color = BLACK;
    }

}
